# 导入需要的库
import gym
import cv2
import numpy as np
from scipy.interpolate import splprep, splev


class DriverEnv(gym.Env):
    # 地图信息
    map = np.nan
    nodes = []
    currentNode = 0
    markColor = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]
    # 车辆信息
    dt = 0.1
    time = 0
    position = np.nan
    velocity = np.nan
    direction = np.nan
    carSize = 6
    # 激励设置
    reward_reachNode = 1000
    reward_clash = -1000
    reward_stay = -0.6
    reward_velocity = 0.1
    # Observation Space
    scans = np.nan
    dir = np.nan
    # Action Space
    actions2acc = np.nan

    def __init__(self, path=None):
        # 加载地图
        if path is None:
            self.mapGenerater()
        else:
            if path.endswith(".png") or path.endswith(".jpg"):
                self.ReadMapFromImg(path)
            else:
                self.readMapAndNodes(path)
        # 初始化动作空间
        accV, accD = 0.3, 0.6
        self.actions2acc = [
            (-accV, -accD * 3),
            (-accV, -accD),
            (-accV, 0),
            (-accV, accD),
            (-accV, accD * 3),
            (0, -accD * 3),
            (0, -accD),
            (0, 0),
            (0, accD),
            (0, accD * 3),
            (accV, accD * 3),
            (accV, accD),
            (accV, 0),
            (accV, -accD),
            (accV, -accD * 3),
            (accV * 3, accD * 3),
            (accV * 3, accD),
            (accV * 3, 0),
            (accV * 3, -accD),
            (accV * 3, -accD * 3),
            (accV * 8, accD * 3),
            (accV * 8, accD),
            (accV * 8, 0),
            (accV * 8, -accD),
            (accV * 8, -accD * 3),
        ]

        self.action_space = gym.spaces.Discrete(len(self.actions2acc))
        # 初始化观测空间,为43维向量，连续值
        lowlim = np.ones(43, dtype=np.float32) * 0
        highlim = np.ones(43, dtype=np.float32) * 100
        lowlim[0] = -100
        lowlim[1] = -np.pi
        highlim[1] = np.pi
        self.observation_space = gym.spaces.Box(lowlim, highlim)

    def reset(self):
        self.position = self.nodes[0].copy()
        self.currentNode = 0
        self.velocity = 0
        v = self.nodes[1] - self.nodes[0]
        self.direction = np.arctan2(v[1], v[0]) + np.random.random() * 0.1 - 0.05
        self.time = 0
        return self.getObservation()

    def step(self, action, k=0):
        # 移动
        self.velocity += self.actions2acc[action][0] * self.dt
        self.direction += self.actions2acc[action][1] * self.dt
        v = self.velocity * np.array(
            [np.cos(self.direction), np.sin(self.direction)], dtype=np.float32
        )
        self.position += v * self.dt
        # 节点更新,这里检测凭据是上一点到当前位置的投影距离大于节点间距,且到下一点距离小于路宽
        done = False
        reward = self.reward_stay  # 为了防止呆在原地，给予一点惩罚
        v1 = self.nodes[self.currentNode + 1] - self.nodes[self.currentNode]
        v2 = self.position - self.nodes[self.currentNode]
        if np.linalg.norm(v1) <= np.dot(v1, v2) and np.linalg.norm(v2 - v1) < 100:
            self.currentNode += 1
            reward += self.reward_reachNode
            if self.currentNode == len(self.nodes) - 1:
                self.currentNode = 0
                done = True
        # 检测撞击
        vertices = self.getVertices()
        for v in vertices:
            if self.map[int(v[1]), int(v[0])] == False:
                reward -= self.reward_clash
                reward -= self.reward_velocity * np.linalg.norm(self.velocity) * (k+0.2)
                done = True
                break
        # 返回
        return self.getObservation(), reward, done, {}

    def render(self):
        img = np.array(self.map, dtype=np.uint8) * 255
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        # 绘制路径
        cv2.polylines(
            img, [np.array(self.nodes, dtype=np.int32)], False, (255, 0, 0), 1
        )
        # 绘制车辆
        vertices = np.array(self.getVertices(), np.int32)
        p1 = vertices[0] * 2 / 3 + vertices[-1] * 1 / 3
        p2 = vertices[0] / 2 + vertices[1] / 2
        p3 = vertices[1] * 2 / 3 + vertices[2] * 1 / 3
        cv2.fillPoly(img, [vertices], color=(0, 255, 0))
        cv2.polylines(img, [np.array([p1, p2, p3], np.int32)], True, color=(0, 0, 255))
        ####################################
        # 绘制扫描结果
        img2 = np.zeros((120, img.shape[1], 3), np.uint8)
        vs = []
        for i in range(len(self.scans)):
            vs.append([i / len(self.scans) * img2.shape[1], self.scans[i]])
        cv2.polylines(img2, [np.array(vs, np.int32)], False, color=(0, 255, 0))
        # 绘制目标方向
        x = int((self.dir + np.pi / 2) / (np.pi) * img.shape[1])
        cv2.line(img2, (x, 0), (x, 120), (255, 255, 0), 2)
        # 合并
        img = np.concatenate([img, img2], axis=0)
        return img

    def getVertices(self):
        c = self.carSize / 2
        vP = np.array([np.cos(self.direction), np.sin(self.direction)])
        vV = np.array(
            [np.cos(self.direction + np.pi / 2), np.sin(self.direction + np.pi / 2)]
        )
        Vertices = []
        Vertices.append(self.position + vP * c * 2 - vV * c)
        Vertices.append(self.position + vP * c * 2 + vV * c)
        Vertices.append(self.position - vP * c * 2 + vV * c)
        Vertices.append(self.position - vP * c * 2 - vV * c)
        return Vertices

    def getObservation(self):
        velo = self.velocity
        v = self.nodes[self.currentNode + 1] - self.position
        dir = np.arctan2(v[1], v[0]) - self.direction
        dir = np.arctan2(np.sin(dir), np.cos(dir))
        scans = self.getScans()
        observation = [[velo], [dir], scans]

        self.scans = scans
        self.dir = dir
        return np.concatenate(observation)

    # 获取当前状态下的扫描信息
    def getScans(self):
        scans = []
        for i in range(-20, 21):
            theta = self.direction + i * np.pi / 40
            scans.append(self.getDist(theta))
        scans = np.array(scans, dtype=np.float32)
        scans = np.clip(scans, 0, 100)
        return scans

    # 获取指定方向的障碍物距离
    def getDist(self, theta):
        p1 = self.position.copy()
        count = 0
        dp = np.array([np.cos(theta), np.sin(theta)])
        while self.map[int(p1[1]), int(p1[0])]:
            p1 += dp
            count += 1
        return count

    # 根据规则解析地图文件
    def ReadMapFromImg(self, path):
        img = cv2.imread(path)
        points = self.findColorPoints(img, self.markColor)
        # 查看起点
        if len(points[0]) == 0:  # 找不到起点就抛异常
            raise Exception("Can't find start point")
        p = points[0][0]
        self.nodes.append(p)
        # 添加中间点,根据距离逐一添加
        while len(points[1]) > 0:
            index, mindist = 0, img.shape[1]
            for i in range(len(points[1])):
                dist = np.sqrt(
                    (p[0] - points[1][i][0]) ** 2 + (p[1] - points[1][i][1]) ** 2
                )
                if dist < mindist:
                    index = i
                    mindist = dist
            p = points[1][index]
            self.nodes.append(p)
            points[1].remove(p)
        # 添加终点
        if len(points[2]) == 0:  # 找不到终点就抛异常
            raise Exception("Can't find end point")
        self.nodes.append(points[2][0])
        self.nodes = np.array(self.nodes, dtype=np.float32)
        # 把地图转化成二维布尔值数组
        self.map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) > 128

    # 在途中寻找指定颜色点
    def findColorPoints(self, img, colors):
        points = [[] for i in range(len(colors))]
        for i in range(img.shape[0]):
            for j in range(img.shape[1]):
                for k in range(len(colors)):
                    if np.array_equal(img[i, j], colors[k]):
                        points[k].append([j, i])
                        img[i, j] = np.array([255, 255, 255], dtype=img.dtype)  # 擦除
        return points

    # 把解析得到的地图和节点保存到文件
    def saveMapAndNodes(self, path):
        n = np.array(self.nodes)
        m = np.array(self.map)
        np.save(path + "_map", m)
        np.save(path + "_nodes", n)

    # 从文件中直接读取地图和节点
    def readMapAndNodes(self, path):
        self.map = np.load(path + "_map.npy")
        self.nodes = np.load(path + "_nodes.npy")

    def mapGenerator(
        self,
        roadBroad=50,
        nodeNum=100,
        nodeInterval_Min=5,
        nodeInterval_Max=40,
        maxTurnAngle=60 / 180 * np.pi,
        pointMaxPosition=[1200, 900],
    ):
        # 初始化参数
        self.nodes = []

        p = np.random.random([2]) * np.array(pointMaxPosition)
        dir = np.pi / 4
        # 生成节点
        self.nodes.append(p.copy())
        while len(self.nodes) < nodeNum:
            dAngle = np.random.random() * maxTurnAngle * 2 - maxTurnAngle
            dir += dAngle
            length = (
                np.random.random() * (nodeInterval_Max - nodeInterval_Min)
                + nodeInterval_Min
            )
            v = np.array([np.cos(dir), np.sin(dir)]) * length
            p += v
            if (
                p[0] < 0
                or p[1] < 0
                or p[0] > pointMaxPosition[0]
                or p[1] > pointMaxPosition[1]
            ):
                p -= v
                continue
            self.nodes.append(p.copy())

        # 节点标准化
        self.nodes = np.array(self.nodes, dtype=np.float32)
        minX, minY = np.min(self.nodes, axis=0)
        self.nodes -= np.array([minX, minY])
        self.nodes += np.array([roadBroad * 2, roadBroad * 2])

        # 使用B样条插值生成平滑曲线
        tck, u = splprep(self.nodes.T, s=0)  # s=0表示不进行平滑处理
        u_new = np.linspace(0, 1, 1000)  # 更多的点使得曲线更平滑
        smooth_nodes = np.array(splev(u_new, tck)).T

        # 确定地图大小
        maxX, maxY = np.max(smooth_nodes, axis=0) + np.array([roadBroad, roadBroad])
        self.map = np.zeros([int(maxY), int(maxX), 3], dtype=np.uint8)

        # 绘制平滑的地图并转化
        cv2.polylines(
            self.map,
            [smooth_nodes.astype(np.int32)],
            False,
            color=(255, 255, 255),
            thickness=roadBroad // 2,
        )
        self.map = cv2.cvtColor(self.map, cv2.COLOR_BGR2GRAY) > 128


if __name__ == "__main__":
    # 初始化
    env = DriverEnv("maps/0")
    env.mapGenerator()
    env.reset()
    # 显示
    action = env.action_space.sample()
    next_state, reward, done, _ = env.step(action)
    img = env.render()
    cv2.imshow("map", img)
    cv2.waitKey(0)